<!DOCTYPE html>
<html>
    <head>
        <meta charset="UTF-8">
        <title>Merge Selector Lab</title>
        <link rel="stylesheet" href="https://stackpath.bootstrapcdn.com/bootstrap/4.5.2/css/bootstrap.min.css">
        <script src="https://d3js.org/d3.v7.min.js"></script>
    </head>
    <body>
        <div class="container-fluid">
            <form id="param-form" class="p-3">
                <div class="row">
                    <div class="form-group col-md-4">
                        <label for="insertPartSize">Initial part size:</label>
                        <input type="text" class="form-control form-control-sm" id="insertPartSize" value="10MB">
                    </div>
                    <div class="form-group col-md-4">
                        <label for="parts">Initial parts count:</label>
                        <input type="number" class="form-control form-control-sm" id="parts" value="243">
                    </div>
                    <div class="form-group col-md-4">
                        <label for="poolSize">Workers in merge pool:</label>
                        <input type="number" class="form-control form-control-sm" id="poolSize" value="9">
                    </div>
                </div>
                <div class="row">
                    <div class="form-group col-md-4">
                        <label for="maxSourceParts">Max parts to merge:</label>
                        <input type="number" class="form-control form-control-sm" id="maxSourceParts" value="100">
                    </div>
                    <div class="form-group col-md-4">
                        <label for="maxPartSize">Average final part size:</label>
                        <input type="text" class="form-control form-control-sm" id="maxPartSize" placeholder="Unlimited">
                    </div>
                    <div class="form-group col-md-4">
                        <label for="method">Method:</label>
                        <select class="form-control form-control-sm" id="method">
                            <!-- Methods will be populated dynamically -->
                        </select>
                    </div>
                    <div class="form-group col-md-8">
                        <label for="layerBases">Bases (number of parts to merge on every level):</label>
                        <input type="text" class="form-control form-control-sm" id="layerBases" value="[3,3,3,3,3]">
                    </div>
                </div>
                <div class="row mt-3">
                    <div class="col">
                        <button type="button" class="btn btn-primary" id="solveButton">Solve</button>
                        <button type="button" class="btn btn-primary mx-2" id="runButton">Run</button>
                    </div>
                </div>
            </form>
        </div>


        <div class="container-fluid">
            <div class="p-3" id="options-container">
                <table class="table table-bordered table-sm" id="options-table">
                    <thead>
                        <tr>
                            <th>WA</th>
                            <th>Integral, parts&times;sec</th>
                            <th>Bases, parts</th>
                            <th></th>
                        </tr>
                    </thead>
                    <tbody></tbody>
                </table>
            </div>
        </div>

        <div class="container-fluid bg-primary text-white">
            <div class="p-3" id="metrics-container"></div>
        </div>

        <div class="container-fluid">
            <div class="row">
                <div class="col-md-12" id="util-container"></div>
            </div>
        </div>

        <div class="container-fluid bg-primary text-white">
            <div class="p-3" id="rewind-container"></div>
        </div>

        <div class="container-fluid">
            <div class="row">
                <div class="col-md-12" id="time-container"></div>
            </div>
        </div>

        <div class="container-fluid">
            <div class="p-3" id="var-container"></div>
        </div>

        <script type="module">
            import { MergeTree } from './MergeTree.js';
            import { SimulationContainer } from './SimulationContainer.js';
            import { sequenceInserter } from './sequenceInserter.js';
            import { customScenario } from './customScenario.js';
            import { floatLayerMerges } from './floatLayerMerges.js';
            import { mergeModel } from './mergeModel.js';
            import { projectedGradientDescent, nesterovAcceleratedPGD, subgradientMethod } from './gradientDescent.js';
            import { ModelSolver } from './ModelSolver.js';
            import { delayMs, parseHumanReadableSize, fetchText } from './util.js';

            function readPartSize() {
                return parseHumanReadableSize(document.getElementById('insertPartSize').value);
            }

            function readMaxPartSize() {
                const str = document.getElementById('maxPartSize').value;
                return str ? parseHumanReadableSize(str) : null; // No limit
            }

            async function loadSolver(modelName) {
                let solver = new ModelSolver(mergeModel);
                let jsonString = await fetchText(`solvers/${modelName}.json`);
                solver.deserializeData(jsonString);
                let result = function(params, onOption) {
                    solver.mergeModel(params, onOption);
                };
                result.isModelSolver = true;
                return result;
            }

            const methods = {
                "Subgradient Method": subgradientMethod,
                "Projected GD": projectedGradientDescent,
                "Accelerated PGD": nesterovAcceleratedPGD,
                "BLI-x4": await loadSolver("solver-64KB_256GB_32768_x4_x4"),
                "BLI-x1": await loadSolver("solver-64KB_256GB_32768"),
            };

            // Populate methods dropdown
            const methodSelect = document.getElementById('method');
            for (const [name] of Object.entries(methods)) {
                const option = document.createElement('option');
                option.value = name;
                option.textContent = name;
                methodSelect.appendChild(option);
            }

            function selectMethod() {
                const selectedMethodName = document.getElementById('method').value;
                return methods[selectedMethodName];
            }

            function readParams() {
                let params = {
                    insertPartSize: readPartSize(),
                    layerBases: JSON.parse(document.getElementById('layerBases').value),
                    parts: parseInt(document.getElementById('parts').value),
                    workers: parseInt(document.getElementById('poolSize').value),
                    maxSourceParts: parseInt(document.getElementById('maxSourceParts').value),
                    maxPartSize: readMaxPartSize(),
                    method: selectMethod()
                };
                console.log("PARAMS", params);
                return params;
            }

            document.getElementById('solveButton').addEventListener('click', async () => {
                const solveButton = document.getElementById('solveButton');
                solveButton.textContent = '⏳';
                solveButton.disabled = true;

                const startTime = performance.now(); // Start timing

                await delayMs(32);

                const {insertPartSize, parts, workers, maxSourceParts, maxPartSize, method} = readParams();

                const mt0 = new MergeTree();
                const realB = mt0.mergeDuration(insertPartSize * parts, parts);

                let Lopt = 1;
                let xopt = [];
                let Fopt = Infinity;

                d3.select('#options-table tbody').html('');

                function solve(params, method, onOption) {
                    if ('isModelSolver' in method) {
                        method(params, onOption);
                    } else {
                        mergeModel(params, method, onOption);
                    }
                }

                solve({insertPartSize, parts, workers, maxSourceParts, maxPartSize},
                    method,
                    function onOption(L, x, Fx) {
                        const bases = x.map(xi => Math.exp(xi));
                        const basesStr = '[' + bases.map(base => base.toFixed(2)).join(', ') + ']';
                        console.log("OPTION", {L, x, Fx, bases});
                        d3.select('#options-table tbody')
                            .append('tr')
                            .attr('data-L', L)
                            .html(`
                                <td>${L + 1}</td>
                                <td>${(Fx * realB).toFixed(2)}</td>
                                <td>${basesStr}</td>
                                <td><button class="btn btn-primary btn-xs choose-btn" style="padding: 2px 5px; font-size: 0.8rem;">Run</button></td>`)
                            .select('button.choose-btn')
                            .on('click', () => {
                                d3.select('#layerBases').property('value', basesStr);
                                document.getElementById('runButton').click();
                            });
                        if (Fx < Fopt) {
                            Lopt = L;
                            xopt = x;
                            Fopt = Fx;
                            d3.select('#layerBases').property('value', basesStr);
                        }
                    }
                );

                // Highlight optimal row
                console.log("OPTIMAL", {Lopt, xopt, Fopt});
                d3.select('#options-table tbody').selectAll('tr')
                    .each(function() {
                        const row = d3.select(this);
                        if (parseInt(row.attr('data-L')) === Lopt) {
                            row.style('background-color', '#d4edda');
                        }
                    });

                const endTime = performance.now(); // End timing
                console.log(`Solve execution time: ${(endTime - startTime).toFixed(2)} ms`); // Log time

                solveButton.textContent = 'Solve';
                solveButton.disabled = false;
            });

            document.getElementById('runButton').addEventListener('click', async () => {
                const {insertPartSize, parts, workers, layerBases} = readParams();

                const scenario = {
                    inserts: [
                        sequenceInserter({
                            start_time: 0,
                            interval: 0.0,
                            parts,
                            bytes: insertPartSize,
                        }),
                    ],
                    selector: floatLayerMerges({insertPartSize, layerBases, mergeRight: true}),
                    pool_size: workers,
                }

                const simContainer = new SimulationContainer();

                async function updateMergeTree({mt}) {
                    simContainer.update({mt}, true);
                    await delayMs(1);
                }

                const signals = {
                    on_every_frame: updateMergeTree,
                };

                const mt = await customScenario(scenario, signals);
                simContainer.update({mt}, true);
                simContainer.setTime(mt.time);
                console.log(mt);
            });
        </script>
    </body>
</html>
